package dan.backend
import  chisel3._
import chisel3.util._
import dan.common._ 
import dan.common.Consts._
import  dan.common.CoreBundle
import  dan.common.UOp
import dan.common.ExpCode
import chisel3.SpecifiedDirection.Flip
import dan.frontend.BrUpdateInfo
import dan.common.TLBCsrContext
import dan.utility.WrapAdd

class Exception extends CoreBundle{
    val uop = new UOp()
    val expCode = Bits(ExpCode.expBits.W)
    val badAddr = UInt(paBits.W)
}

class LSUDMemIO extends CoreBundle{
    val robBits = robParam.robIdxBits
    val req = new DecoupledIO(Vec(memWidth, Valid(new DcacheReq)))
    val s1Kill = Output(Vec(memWidth, Bool()))
    val s2Hit = Input(Vec(memWidth, Bool()))
    val resp = Flipped(Vec(memWidth, new Valid(new DcacheResp)))
    /* 
    Negative Acknowledgement (拒绝服务) 信号。这是一种比 req.ready 更高级的流控机制。
    有时候D-Cache因为内部资源冲突（如MSHR已满、写缓冲已满）暂时无法处理某个请求，但又不想阻塞整个LSU流水线。
    此时，D-Cache可以先接收请求，然后立即通过 nack 通道将原始请求发回给LSU，并附上 valid 信号。LSU收到 nack 后，需要在稍后的时间点重新发送这个请求。
     */
    val negAck = Flipped(Vec(memWidth, new Valid(new DcacheReq)))
    /* 
        分支预测旁路优化
     */
    val brUpdate = Output(new BrUpdateInfo())
    val exception = Output(Bool())
    val isLL = Output(Bool())
    val memFence = Output(Bool())
    /* 
        命令LSU必须停止发送新的内存请求，直到该信号变低
     */
    val ordered = Input(Bool())
}

class LSUExeIO extends CoreBundle {
    val req   = Flipped(new Valid(new FuncUnitResp))
    val iresp = new DecoupledIO(new ExecUnitResp)
}

class LSUTLBDataIO extends CoreBundle {
    val r_req  = Vec(memWidth, Valid(new TLBDataRReq))
    val r_resp = Vec(memWidth, Flipped(Valid(new TLBDataRResp)))
}

class LSUCsrIO extends CoreBundle {
    val dtlb_csr_ctx = Input(new TLBCsrContext)
    val llbit        = Input(Bool())
}

/**
 * 输入dispatch阶段的uop，commit的信息，rob，brupdate
 * 输出ldq，stq的索引，是否full，给rob的clear信号
 */
// lsu <> core
class LSUCoreIO extends CoreBundle {
    /** *********************************** */
    val robAddrSz = robParam.robIdxBits
    val ldqAddrSz = lsuParam.ldqIdxBits
    val stqAddrSz = lsuParam.stqIdxBits
    /** ************************************ */

    val exe = Vec(memWidth, new LSUExeIO)

    val dis_uops    = Flipped(Vec(coreWidth, Valid(new UOp)))
    val dis_ldq_idx = Output(Vec(coreWidth, UInt(ldqAddrSz.W)))
    val dis_stq_idx = Output(Vec(coreWidth, UInt(stqAddrSz.W)))

    val ldq_full    = Output(Vec(coreWidth,Bool()))
    val stq_full    = Output(Vec(coreWidth,Bool()))

    val commit      = Input(new CommitSignal)
    val commit_load_at_rob_head = Input(Bool())

    val clr_bsy     = Output(Vec(memWidth,Valid(UInt(robAddrSz.W))))    // TODO

    val fence_dmem  = Input(Bool())

    // Speculatively tell the IQs that we'll get load data back next cycle
    val spec_ld_wakeup = Output(Vec(memWidth, Valid(UInt(physRegBits.W))))
    val ld_miss     = Output(Bool())

    val brupdate    = Input(new BrUpdateInfo)
    val rob_head_idx= Input(UInt(robAddrSz.W))
    val exception   = Input(Bool())

    val stq_empty   = Output(Bool())
    val dcache_ord  = Output(Bool())

    val lsu_xcpt    = Output(Valid(new Exception))

    val tlb_data    = new LSUTLBDataIO
}

class LDQEntry extends CoreBundle with HasUOp {
    /** *********************************** */
    val numStqEntries = lsuParam.stqSize
    val numLdqEntries = lsuParam.ldqSize
    val stqAddrSz     = lsuParam.stqIdxBits
    val ldqAddrSz     = lsuParam.ldqIdxBits

    // require(isPow2(numLdqEntries), "numLdqEntries must be a power of 2")
    // require(isPow2(numStqEntries), "numStqEntries must be a power of 2")
    /** ************************************ */

    val addr             = Valid(UInt(xBits.W))
    val xcpt_valid       = Bool()

    val executed         = Bool()
    val succeeded        = Bool() // dcache 返回结果
    val order_fail       = Bool() // raw 冒险

    val st_dep_mask      = UInt(numStqEntries.W)
    val youngest_stq_idx = UInt(stqAddrSz.W) // 第一条晚于该load的store指令

    val forward_std_val  = Bool()
    val forward_stq_idx  = UInt(stqAddrSz.W)

    val is_uncacheable   = Bool()
}

class STQEntry extends CoreBundle with HasUOp {
    val addr            = Valid(UInt(xBits.W))
    val xcpt_valid      = Bool()
    val data            = Valid(UInt(xBits.W))
    val is_uncacheable  = Bool()

    val committed       = Bool()
    val succeeded       = Bool() // 访存成功
}

//////////////////////////////
class IndexAllocator extends CoreModule {
    val stqAddrSz       = lsuParam.stqIdxBits
    val ldqAddrSz       = lsuParam.ldqIdxBits
    val numStqEntries   = lsuParam.stqSize
    val numLdqEntries   = lsuParam.ldqSize
    val io = IO(new CoreBundle {
        val dis_uops    = Input(Vec(coreWidth, Valid(new UOp)))
        val old_ldq_tail = Input(UInt(ldqAddrSz.W))
        val old_stq_tail = Input(UInt(stqAddrSz.W))
        val ldq_enq_idxs = Output(Vec(coreWidth, UInt(ldqAddrSz.W)))
        val stq_enq_idxs = Output(Vec(coreWidth, UInt(stqAddrSz.W)))
        val new_ldq_tail     = Output(UInt(ldqAddrSz.W))
        val new_stq_tail     = Output(UInt(stqAddrSz.W))
    })
    val dis_ld_valids = VecInit(io.dis_uops.map(x=>x.valid && x.bits.useLdQ && !x.bits.expValid))
    val dis_st_valids = VecInit(io.dis_uops.map(x=>x.valid && x.bits.useStQ && !x.bits.expValid))
    // assert(coreWidth == 4 , "coreWidth must be 4")
    io.ldq_enq_idxs := 0.U.asTypeOf(Vec(coreWidth, UInt(ldqAddrSz.W)))
    io.stq_enq_idxs := 0.U.asTypeOf(Vec(coreWidth, UInt(stqAddrSz.W)))
    io.new_ldq_tail := 0.U
    io.new_stq_tail := 0.U

    val ldq_tail_add_1 = WrapAdd(io.old_ldq_tail, 1.U, numLdqEntries)
    val ldq_tail_add_2 = WrapAdd(io.old_ldq_tail, 2.U, numLdqEntries)
    val ldq_tail_add_3 = WrapAdd(io.old_ldq_tail, 3.U, numLdqEntries)

    val stq_tail_add_1 = WrapAdd(io.old_stq_tail, 1.U, numStqEntries)
    val stq_tail_add_2 = WrapAdd(io.old_stq_tail, 2.U, numStqEntries)
    val stq_tail_add_3 = WrapAdd(io.old_stq_tail, 3.U, numStqEntries)

    assert(coreWidth == 3 , "coreWidth must be 3")
    switch(dis_ld_valids.asUInt){
        is("b000".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := io.old_ldq_tail
            io.ldq_enq_idxs(2) := io.old_ldq_tail
            io.new_ldq_tail := io.old_ldq_tail
        }
        is("b001".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := ldq_tail_add_1
            io.ldq_enq_idxs(2) := ldq_tail_add_1
            io.new_ldq_tail := ldq_tail_add_1
        }
        is("b010".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := io.old_ldq_tail
            io.ldq_enq_idxs(2) := ldq_tail_add_1
            io.new_ldq_tail := ldq_tail_add_1
        }
        is("b011".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := ldq_tail_add_1
            io.ldq_enq_idxs(2) := ldq_tail_add_2
            io.new_ldq_tail := ldq_tail_add_2
        }
        is("b100".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := io.old_ldq_tail
            io.ldq_enq_idxs(2) := io.old_ldq_tail
            io.new_ldq_tail := ldq_tail_add_1
        }
        is("b101".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := ldq_tail_add_1
            io.ldq_enq_idxs(2) := ldq_tail_add_1
            io.new_ldq_tail := ldq_tail_add_2
        }
        is("b110".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := io.old_ldq_tail
            io.ldq_enq_idxs(2) := ldq_tail_add_1
            io.new_ldq_tail := ldq_tail_add_2
        }
        is("b111".U){
            io.ldq_enq_idxs(0) := io.old_ldq_tail
            io.ldq_enq_idxs(1) := ldq_tail_add_1
            io.ldq_enq_idxs(2) := ldq_tail_add_2
            io.new_ldq_tail := ldq_tail_add_3
        }
    }

    switch(dis_st_valids.asUInt){
        is("b000".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := io.old_stq_tail
            io.stq_enq_idxs(2) := io.old_stq_tail
            io.new_stq_tail := io.old_stq_tail
        }
        is("b001".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := stq_tail_add_1
            io.stq_enq_idxs(2) := stq_tail_add_1
            io.new_stq_tail := stq_tail_add_1
        }
        is("b010".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := io.old_stq_tail
            io.stq_enq_idxs(2) := stq_tail_add_1
            io.new_stq_tail := stq_tail_add_1
        }
        is("b011".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := stq_tail_add_1
            io.stq_enq_idxs(2) := stq_tail_add_2
            io.new_stq_tail := stq_tail_add_2
        }
        is("b100".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := io.old_stq_tail
            io.stq_enq_idxs(2) := io.old_stq_tail
            io.new_stq_tail := stq_tail_add_1
        }
        is("b101".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := stq_tail_add_1
            io.stq_enq_idxs(2) := stq_tail_add_1
            io.new_stq_tail := stq_tail_add_2
        }
        is("b110".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := io.old_stq_tail
            io.stq_enq_idxs(2) := stq_tail_add_1
            io.new_stq_tail := stq_tail_add_2
        }
        is("b111".U){
            io.stq_enq_idxs(0) := io.old_stq_tail
            io.stq_enq_idxs(1) := stq_tail_add_1
            io.stq_enq_idxs(2) := stq_tail_add_2
            io.new_stq_tail := stq_tail_add_3
        }
    }
}
